package bloomfilter

import (
	"hash/fnv"
	"math"
	"math/rand"
)

type BloomFilter struct {
	// 位图
	bitmap []bool
	// 位图中 bit 长度
	size uint
	// 每个哈希函数对应一个不同种子
	seeds []uint64
}

func NewBloomFilter(expectedItems uint, falsePositiveRate float64) *BloomFilter {
	size, numHashFunc := calculateParameters(expectedItems, falsePositiveRate)

	// 随机生成不同的种子
	seeds := make([]uint64, numHashFunc)
	for i := range seeds {
		seeds[i] = rand.Uint64()
	}

	return &BloomFilter{
		bitmap: make([]bool, size),
		size:   size,
		seeds:  seeds,
	}
}

// 计算哈希值（注入种子）
func (bf *BloomFilter) getHashes(item []byte) []uint64 {
	hashes := make([]uint64, len(bf.seeds))
	h := fnv.New64a()
	for i, seed := range bf.seeds {
		h.Reset()
		h.Write(item)
		h.Write([]byte{byte(seed)}) // 注入种子
		hashes[i] = h.Sum64()
	}
	return hashes
}

func (bf *BloomFilter) Add(item []byte) {
	for _, hash := range bf.getHashes(item) {
		bf.bitmap[hash%uint64(bf.size)] = true
	}
}

func (bf *BloomFilter) Contains(item []byte) bool {
	for _, hash := range bf.getHashes(item) {
		if !bf.bitmap[hash%uint64(bf.size)] {
			return false
		}
	}

	return true
}

type BloomFilterV1 struct {
	// 位图
	bitmap []byte
	// 位图中 bit 长度
	size uint
	// 每个哈希函数对应一个不同种子，切片长度就是哈希函数的个数
	seeds []uint64
}

func NewBloomFilterV1(expectedItems uint, falsePositiveRate float64) *BloomFilterV1 {
	size, numHashFunc := calculateParameters(expectedItems, falsePositiveRate)

	// 随机生成不同的种子
	seeds := make([]uint64, numHashFunc)
	for i := range seeds {
		seeds[i] = rand.Uint64()
	}

	return &BloomFilterV1{
		bitmap: make([]byte, (size+7)/8), // byte 数量向上取整
		size:   size,
		seeds:  seeds,
	}
}

// 计算哈希值
func (bf *BloomFilterV1) Add(item []byte) {
	h := fnv.New64a()
	for _, seed := range bf.seeds {
		h.Reset()
		h.Write(item)
		h.Write([]byte{byte(seed)}) // 注入种子
		index := h.Sum64() % uint64(bf.size)
		bf.setBit(index)
	}
}

func (bf *BloomFilterV1) setBit(index uint64) {
	bf.bitmap[index/8] |= (1 << (index % 8))
}

func (bf *BloomFilterV1) Contains(item []byte) bool {
	h := fnv.New64a()
	for _, seed := range bf.seeds {
		h.Reset()
		h.Write(item)
		h.Write([]byte{byte(seed)}) // 注入种子
		index := h.Sum64() % uint64(bf.size)
		if !bf.getBit(index) {
			return false
		}
	}
	return true
}

func (bf *BloomFilterV1) getBit(index uint64) bool {
	return bf.bitmap[index/8]&(1<<(index%8)) > 0
}

// 计算最优参数
// n: 预计存储的数据量，
// p: 预期误差率 0~1，
// size: bitmap 长度
// numHashFunc: hash 函数个数
func calculateParameters(n uint, p float64) (size uint, numHashFunc uint) {
	size = uint(math.Ceil(-(float64(n) * math.Log(p)) / (math.Log(2) * math.Log(2))))
	numHashFunc = uint(math.Ceil(math.Log(2) * float64(size) / float64(n)))
	return
}
